/*
 * CloudStego school project
 * @author lauris kruusamae
 * @author triin uudam
 * @author uku loskit
 */
package main;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.CommandLineParser;
import org.apache.commons.cli.GnuParser;
import org.apache.commons.cli.Options;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.IntWritable;
import org.apache.hadoop.io.SequenceFile;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.lib.input.FileInputFormat;
import org.apache.hadoop.mapreduce.lib.input.SequenceFileInputFormat;
import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat;
import org.apache.hadoop.mapreduce.lib.output.SequenceFileOutputFormat;

/**
 * Main class.
 * Reads CLI arguments
 * Sets up the job configuration
 * Starts the job 
 *
 */
public class MapReduceSkeleton {
    
    
    static Configuration conf;
    
    
    /**
     * Runs the Hadoop MapReduce job
     * @param pathin
     * @param pathout
     * @param nrOfTasks
     * @param encode
     * @return
     * @throws Exception
     */
    public int run(String pathin, String pathout, int nrOfTasks, boolean encode) throws Exception {
        Job job = new Job(conf);
        
        job.setJarByClass(getClass());

        if (encode) {
        	
        	// Mapper and Reducer for encoding
        	job.setMapperClass(EncodeMap.class);
        	job.setReducerClass(EncodeReduce.class);
        } else {
        	
        	// Mapper and reducer for decoding
        	job.setMapperClass(DecodeMap.class);
        	job.setReducerClass(DecodeReduce.class);
        }
        
        job.setInputFormatClass(SequenceFileInputFormat.class);
        job.setOutputFormatClass(SequenceFileOutputFormat.class);
        
        job.setMapOutputKeyClass(Text.class);
        job.setMapOutputValueClass(Text.class);
        
        job.setOutputKeyClass(IntWritable.class);
        job.setOutputValueClass(IntWritable.class);    
        
        job.setNumReduceTasks(nrOfTasks);
        
        Path outPath = new Path(pathout);
        
        FileSystem fs = FileSystem.get(conf);
        if (fs.exists(outPath)) {
        	fs.delete(outPath, true);
        }
        
        FileInputFormat.addInputPath(job, new Path(pathin));
        FileOutputFormat.setOutputPath(job, outPath);

        boolean success = job.waitForCompletion(true);
        return success ? 0 : -1;
    }
    

    public static void main(String[] args) throws Exception {
        
        MapReduceSkeleton skel = new MapReduceSkeleton();
        
        conf = new Configuration();
        
        // create Options object
        Options options = new Options();
        options.addOption("encode", false, "encode");
        options.addOption("decode", false, "decode");
        options.addOption("baseDir", true, "directory where base image files are located");
        options.addOption("input", true, "message image");
        options.addOption("outDir", true, "output directory");
        
        CommandLineParser parser = new GnuParser();
        CommandLine cmd = parser.parse(options, args);
        
        boolean encode = cmd.hasOption("encode");
        boolean decode = cmd.hasOption("decode");
        if (encode && decode || !encode && !decode) {
        	System.out.println("encode or decode option must be given, but not both at the same time");
        	System.exit(0);
        }
        String input = cmd.getOptionValue("input");
        String baseDir = cmd.getOptionValue("baseDir");
        String outputDir = cmd.getOptionValue("outDir");
        
        if (encode) {
        	
        	// Start message encoding process
        	if (input == null || baseDir == null || outputDir == null) {
        		System.out.println("input, baseDir and outDir must be given");
        		System.exit(0);
        	}
        	
        	List<InputData> inputs = readBaseFiles(baseDir, input);
        	if (inputs.isEmpty()) {
        		System.out.println("No png files found in base directory");
        		System.exit(0);
        	}
        	
        	Path outPath = new Path(outputDir);
        	String sequencePath = outPath.getParent().toString() + "sequenceFile";
        	generateSequenceFileForEncode(inputs, sequencePath);  
        	
        	conf.set(EncodeReduce.KEY_OUTPUT_DIR_PATH, outPath.toString());
        	
        	// Since we already checked if encode/decode exist then passing the value of encode is enough
        	skel.run(sequencePath, outputDir, inputs.size(), true);
        } else {
        	
        	// Start decoding process
        	if (baseDir == null || outputDir == null) {
        		System.out.println("baseDir and outDir must be given");
        		System.exit(0);
        	}
        	
        	List<String> inputs = readBaseFilesForDecode(baseDir);
        	if (inputs.isEmpty()) {
        		System.out.println("No png files found in base directory");
        		System.exit(0);
        	}
        	
        	Path outPath = new Path(outputDir);
        	String sequencePath = outPath.getParent().toString() + "sequenceFile";
        	generateSequenceFileForDecode(inputs, sequencePath);  
        	
        	conf.set(EncodeReduce.KEY_OUTPUT_DIR_PATH, outPath.toString());
        	
        	// Since we already checked if encode/decode exist then passing the value of encode is enough
        	skel.run(sequencePath, outputDir, inputs.size(), false);
        }
    }
    
    /**
     * Method to write a key - value pairs to sequence file
     * key - base image path
     * value - data image path
     */
    public static void generateSequenceFileForEncode(List<InputData> inputs, String sequencePath) throws IOException {
    	FileSystem fs = FileSystem.get(conf);
    	Path path = new Path(sequencePath);
    	SequenceFile.Writer writer = SequenceFile.createWriter(fs, conf, path, 
    			Text.class, Text.class, 
    			SequenceFile.CompressionType.NONE);

    	for (InputData input : inputs) {
    		writer.append(new Text(input.getBaseImagePath()), new Text(input.getDataImagePath()));
    	}
    	writer.close();
    }
    
    /**
     * Method to write a key - value pairs to sequence file
     * key - file name
     * value - file path
     */
    public static void generateSequenceFileForDecode(List<String> inputs, String sequencePath) throws IOException {
    	FileSystem fs = FileSystem.get(conf);
    	Path path = new Path(sequencePath);
    	SequenceFile.Writer writer = SequenceFile.createWriter(fs, conf, path, 
    			Text.class, Text.class, 
    			SequenceFile.CompressionType.NONE);

    	for (String pathString : inputs) {
    		writer.append(new Text(Util.getFileName(pathString)), new Text(pathString));
    	}
    	writer.close();
    }
    
    /**
     * Creates InputData objects of all png files in the given base directory
     * @param baseDirPath
     * @param inputFilePath
     * @return
     * @throws IOException
     */
    public static List<InputData> readBaseFiles(String baseDirPath, String inputFilePath) throws IOException {
    	List<InputData> result = new ArrayList<InputData>();
    	FileSystem fs = FileSystem.get(conf);
        Path dir = new Path(baseDirPath);
        FileStatus[] stats = fs.listStatus(dir);
        for (FileStatus stat : stats) {
        	if (stat.isDir()) {
        		continue;
        	}
        	String path = stat.getPath().toUri().getPath();
        	String extension = Util.getFileExtension(path);
        	if (extension != null && extension.equalsIgnoreCase("png")) {
        		// This is a supported file type, add it
        		result.add(new InputData(path, inputFilePath));
        	}
        } 
        return result;
    }
    
    /**
     * Returns a List of path Strings of all png files in the given base directory
     * @param baseDirPath
     * @return
     * @throws IOException
     */
    public static List<String> readBaseFilesForDecode(String baseDirPath) throws IOException {
    	List<String> result = new ArrayList<String>();
    	FileSystem fs = FileSystem.get(conf);
        Path dir = new Path(baseDirPath);
        FileStatus[] stats = fs.listStatus(dir);
        for (FileStatus stat : stats) {
        	if (stat.isDir()) {
        		continue;
        	}
        	String path = stat.getPath().toUri().getPath();
        	String extension = Util.getFileExtension(path);
        	if (extension != null && extension.equalsIgnoreCase("png")) {
        		// This is a supported file type, add it
        		result.add(path);
        	}
        } 
        return result;
    }
}


